import numpy as np
import matplotlib.pyplot as plt
from gher.relabeler.relabeler import RandomRelabeler


class USVRandomRelabeler(RandomRelabeler):
    def __init__(self, path='line', **kwargs):
        super().__init__(**kwargs)
        self.path = path
        self.latent = self.sample_task()

    def calculate_path_reward(self, path, latent):
        pass

    def get_reward_matrix(self, paths, latents):
        pass

    def reward_done(self, obs, action, latent, env_infos=None):
        pass

    def sample_task(self):
        task = None
        if self.path == 'line':
            # 直线路径初始化
            x_0 = np.random.uniform(low=-2.5, high=2.5)
            y_0 = np.random.uniform(low=-2.5, high=2.5)
            x_d = np.random.uniform(low=15, high=30)
            y_d = np.sqrt(100 - (x_0 - x_d) ** 2) + y_0

        elif self.path == 'sin':
            # 正弦轨迹初始化
            angle = np.random.uniform(-np.pi, np.pi, 1)[0]
            d = np.random.uniform(0.75, 1, 1)[0]  # used to be 0.5
            a = np.random.uniform(-0.25, 0.25, 1)[0]
            task = np.array([angle, d, a])
        elif self.path == 'circle':
            # 圆形轨迹初始化
            pass
        return task



    def plot_trajectory_on_heatmap(self, latent, path, title):
        locs = path['observations']

        dx, dy = 0.01, 0.01
        y, x = np.mgrid[slice(-1, 1 + dy, dy),
                        slice(-1, 1 + dx, dx)]
        mesh_xs = np.stack([x, y], axis=2).reshape(-1, 2)
        # ll = torch.exp(self.made.get_log_prob(torch.tensor(mesh_xs).to(device).float())).cpu().detach().numpy()
        # raise NotImplementedError
        rewards = self.calculate_path_reward(dict(observations=mesh_xs), latent)
        fig = plt.figure()
        ax = fig.add_subplot(111)
        c = ax.pcolor(x, y, rewards.reshape([y.shape[0], y.shape[1]]))
        ax.plot(locs[:, 0], locs[:, 1], c='r')
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_title("Reward Heatmap")
        ax.set_aspect('equal')
        fig.colorbar(c, ax=ax)
        plt.savefig('/tmp/heatmap_traj/{}.png'.format(title))
        plt.close('all')

# class USVRelabelerWithObsAndGoal():
#     def __init__(self, sparse_reward=False, **kwargs):
#         su
